package org.gzc.tree;


import java.util.*;

/**
 * @Description: 哈夫曼编码的实现
 * @Author: guozongchao
 * @Date: 2020/8/13 18:56
 */
public class HuffmanCode {
    //存放每个叶子结点对应的哈夫曼编码，即哈夫曼编码表
    private Map<Byte, String> huffmanCodeMap = new HashMap<>();

    //创建结点，数据和权值
    private class Node implements Comparable<Node> {
        Byte data; //存放数据(字符)本身，比如 'a'=>97  , ' ' =>32
        int weight; //权值，表示字符出现的次数
        Node left;
        Node right;

        public Node(Byte data, int weight) {
            this.data = data;
            this.weight = weight;
        }

        @Override
        public int compareTo(Node o) {
            //从小到大排序
            return this.weight - o.weight;
        }

        @Override
        public String toString() {
            return "Node{" +
                    "data=" + data +
                    ", weight=" + weight + '}';
        }

        //前序遍历
        public void preOrder() {
            System.out.println(this);
            if (this.left != null) {
                this.left.preOrder();
            }
            if (this.right != null) {
                this.right.preOrder();
            }
        }
    }

    /**
     * 根据原始字节数组，统计统计字节数组中的某个字符出现的频率次数，构建成一个带权值的结点
     * 返回所有结点的列表集合，用于后续构建 哈夫曼树
     *
     * @param bytes 原始字节数组
     * @return
     */
    private List<Node> getNodes(byte[] bytes) {

        //1、创建一个ArrayList
        List<Node> nodes = new ArrayList<>();

        //遍历bytes，统计每一个byte出现的次数
        //通过map存储起来
        Map<Byte, Integer> counts = new HashMap<>();
        for (byte b : bytes) {
            Integer count = counts.get(b);
            if (count == null) {
                counts.put(b, 1);
            } else {
                counts.put(b, count + 1);
            }
        }

        //将每一个键值对转成一个Node对象，并加入到nodes列表中
        for (Map.Entry<Byte, Integer> entry : counts.entrySet()) {
            nodes.add(new Node(entry.getKey(), entry.getValue()));
        }

        return nodes;
    }

    /**
     * 根据已经创建的 带权值结点的集合，构建 一个 哈夫曼树
     * 并返回该树的根结点
     *
     * @param nodes 带权值结点列表
     * @return
     */
    private Node createHuffmanTree(List<Node> nodes) {
        while (nodes.size() > 1) {
            //排序 ，从小到大
            Collections.sort(nodes);
            //取出第一棵最小的二叉树
            Node leftNode = nodes.get(0);
            //取出第二棵最小的二叉树
            Node rightNode = nodes.get(1);

            //创建一棵新的二叉树，它的根结点没有 data，只有权值
            Node parent = new Node(null, leftNode.weight + rightNode.weight);
            parent.left = leftNode;
            parent.right = rightNode;

            //从nodes中删除已经使用过的结点leftNode 和 rightNode
            //将已经创建好的二叉树的根结点放入nodes中
            nodes.remove(leftNode);
            nodes.remove(rightNode);
            nodes.add(parent);
        }
        return nodes.get(0);
    }

    /**
     * 获得哈夫曼树的叶子结点的哈夫曼编码，保存在一个 huffmanCodeMap 中
     *
     * @param node          哈夫曼树
     * @param code          哈夫曼树的分支编码，例如 左分支为 0，右分支为 1
     * @param stringBuilder 拼接叶子结点的编码
     */
    private void getHuffmanCode(Node node, String code, StringBuilder stringBuilder) {
        StringBuilder sb = new StringBuilder(stringBuilder);
        sb.append(code);
        if (node != null) {
            if (node.data == null) { //如果该结点data为空，说明该结点不是叶子结点
                //递归左结点
                getHuffmanCode(node.left, "0", sb);
                //递归右结点
                getHuffmanCode(node.right, "1", sb);
            } else {  //否则，说明是叶子结点
                //将该结点的原始字符的Byte 和 最终的编码 加入到 map中
                huffmanCodeMap.put(node.data, sb.toString());
            }
        }
    }

    /**
     * 通过哈夫曼树的根结点，获取该树的所有叶子结点的编码
     * 保存 到 huffmanCodeMap 中
     *
     * @param root 哈夫曼树的根结点
     * @return huffmanCodeMap
     */
    private Map<Byte, String> getHuffmanCodeMap(Node root) {
        if (root == null) {
            return null;
        }
        //用于拼接，需要到达叶子结点路径，如010...
        StringBuilder sb = new StringBuilder();
        //处理左子树
        getHuffmanCode(root.left, "0", sb);
        //处理右子树
        getHuffmanCode(root.right, "1", sb);

        //返回哈夫曼编码表
        return huffmanCodeMap;
    }

    /**
     * 通过 原始数据的字节数组 和 对应的哈夫曼编码表 获得编码后的 字节数组
     *
     * @param originalBytes 原数数据的字节数组
     * @return 返回编码后的字节数组
     */
    private byte[] encode(byte[] originalBytes, Map<Byte, String> huffmanCodeMap) {
        //1.利用 huffmanCodeMap 将原始 originalBytes 转成 赫夫曼编码对应的字符串
        StringBuilder sb = new StringBuilder(); //用于拼接
        for (byte b : originalBytes) {
            sb.append(huffmanCodeMap.get(b));
        }

        System.out.println("编码后字符串："+sb.toString());
        //2.将得到的字符串编码 转成字节数组，即每8位为一个字节
        //注意：不是调用getBytes方法
        //2.1  计算字节数组需要多少空间
        int len = sb.length() % 8 == 0 ? sb.length() / 8 : sb.length() / 8 + 1;
        byte[] encodeBytes = new byte[len];
        //2.2  转换成字节数组
        int index = 0;
        for (int i = 0; i < sb.length(); i += 8) {
            String str;
            //防止越界
            if (i + 8 > sb.length()) {  //不够8位
                str = sb.substring(i);
            } else {
                str = sb.substring(i, i + 8);
            }
            //将8位二进制的字符串 按照二进制形式，解析成 整型 ，例如 parseInt("1100110", 2) returns 102
            encodeBytes[index++] = (byte) Integer.parseInt(str, 2);
        }
        return encodeBytes;
    }

    /**
     * 编码
     * 为了方便调用，封装所有方法，直接传入原始字节数组，返回 编码后字节数组
     *
     * @param originalBytes 原数字节数组
     * @return
     */
    public byte[] encode(byte[] originalBytes) {
        List<Node> nodes = getNodes(originalBytes);  //获得需要创建成哈夫曼树的所有带权叶子节点集合
        Node huffmanTree = createHuffmanTree(nodes); //创建哈夫曼树，返回根结点
        Map<Byte, String> huffmanCodeMap = getHuffmanCodeMap(huffmanTree); // 根据哈夫曼树返回 哈夫曼编码表
        byte[] encodeBytes = encode(originalBytes, huffmanCodeMap);  //根据编码表对原来字节数组进行编码，返回编码后字节数组
        return encodeBytes;
    }

    /**
     * 将一个 byte 转成一个二进制的字符串
     * @param b 传入的字节
     * @param flag  标志是否需要补高位如果是 true ，表示需要补高位，如果是 false 表示不补, 如果是最后一个 字节，无需补高位
     * @return  返回传入字节的二进制字符串
     */
    private String byteToBitString(byte b, boolean flag) {
        //使用变量保存 b
        int temp = b;    //将 b 转成 int
        // 如果是正数我们还存在补高位
        if (flag) {
            temp |= 256;   //按位或 256 1 0000 0000 | 0000 0001 => 1 0000 0001
        }
        String str = Integer.toBinaryString(temp); //返回的是 temp 对应的二进制的补码
        if (flag) {
            return str.substring(str.length() - 8);
        } else {
            return str;
        }
    }

    /**
     * 通过 哈夫曼编码表 对已编码的字节数组 解码
     * @param huffmanCodeMap 哈夫曼编码表
     * @param encodeBytes 编码的字节数组
     * @return 返回解码后的字节数组
     */
    private byte[] decode(Map<Byte, String> huffmanCodeMap, byte[] encodeBytes) {
        System.out.println(huffmanCodeMap);
        //1. 先得到 huffmanBytes 对应的 二进制的字符串 ， 形式 1010100010111...
        StringBuilder stringBuilder = new StringBuilder();
        // 将 byte 数组转成二进制的字符串
        for (int i = 0; i < encodeBytes.length; i++) {
            byte b = encodeBytes[i]; //判断是不是最后一个字节
            boolean flag = (i == encodeBytes.length - 1);
            stringBuilder.append(byteToBitString(b, !flag));
        }
        //把字符串安装指定的赫夫曼编码进行解码
        //把赫夫曼编码表进行调换，因为反向查询 a->100 100->a
        Map<String, Byte> map = new HashMap<String, Byte>();
        for (Map.Entry<Byte, String> entry : huffmanCodeMap.entrySet()) {
            map.put(entry.getValue(), entry.getKey());
        }
        //创建要给集合，存放 byte
        List<Byte> list = new ArrayList<>();
        // i 可以理解成就是索引,扫描 stringBuilder
        for (int i = 0; i < stringBuilder.length(); ) {
            int count = 1; // 小的计数器
            boolean flag = true;
            Byte b = null;
            while (flag) {
                //1010100010111...
                // 递增的取出 key 1
                String key= stringBuilder.substring(i, i + count); //i 不动，让 count 移动，指定匹配到一个字符

                b = map.get(key);
                if (b == null) {  //说明没有匹配到
                    count++;
                } else {//匹配到
                    flag = false;
                }
            }
            list.add(b);
            i += count;//i 直接移动到 count
        }
        //当 for 循环结束后，我们 list 中就存放了所有的字符 "i like like like java do you like a java"
        // 把 list 中的数据放入到 byte[] 并返回
        byte b[] = new byte[list.size()];
        for (int i = 0; i < b.length; i++) {
            b[i] = list.get(i);
        }
        return b;
    }

    /**
     * 解码
     * @param encodeBytes  已编码字节数组
     * @return
     */
    public byte[] decode(byte[] encodeBytes){
        byte[] decodeBytes = decode(huffmanCodeMap, encodeBytes);
        return decodeBytes;
    }
}

